import os
import json
import numpy as np
import logging
from typing import List, Dict, Tuple
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import pickle
import re

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RAGSystem:
    """检索增强生成系统"""
    
    def __init__(self, model_name='all-MiniLM-L6-v2', knowledge_base_path='knowledge_base'):
        """
        初始化RAG系统
        
        Args:
            model_name: 句子嵌入模型名称
            knowledge_base_path: 知识库存储路径
        """
        self.model_name = model_name
        self.knowledge_base_path = knowledge_base_path
        self.embedding_model = None
        self.knowledge_base = []
        self.embeddings = None
        self.knowledge_file = os.path.join(knowledge_base_path, 'knowledge_base.json')
        self.embeddings_file = os.path.join(knowledge_base_path, 'embeddings.pkl')
        
        # 创建知识库目录
        os.makedirs(knowledge_base_path, exist_ok=True)
        
        # 初始化嵌入模型
        self._load_embedding_model()
        
        # 加载或创建知识库
        self._load_knowledge_base()
    
    def _load_embedding_model(self):
        """加载句子嵌入模型"""
        try:
            logger.info(f"加载嵌入模型: {self.model_name}")
            self.embedding_model = SentenceTransformer(self.model_name)
            logger.info("嵌入模型加载成功")
        except Exception as e:
            logger.error(f"加载嵌入模型失败: {e}")
            raise
    
    def _load_knowledge_base(self):
        """加载知识库"""
        try:
            if os.path.exists(self.knowledge_file) and os.path.exists(self.embeddings_file):
                logger.info("加载现有知识库...")
                with open(self.knowledge_file, 'r', encoding='utf-8') as f:
                    self.knowledge_base = json.load(f)
                
                with open(self.embeddings_file, 'rb') as f:
                    self.embeddings = pickle.load(f)
                
                logger.info(f"知识库加载完成，包含 {len(self.knowledge_base)} 条知识")
            else:
                logger.info("创建新的知识库...")
                self._create_default_knowledge_base()
        except Exception as e:
            logger.error(f"加载知识库失败: {e}")
            self._create_default_knowledge_base()
    
    def _create_default_knowledge_base(self):
        """创建默认的乳腺癌知识库"""
        default_knowledge = [
            {
                "id": 1,
                "title": "乳腺癌概述",
                "content": "乳腺癌是女性最常见的恶性肿瘤之一，起源于乳腺上皮细胞。早期发现和治疗可以显著提高治愈率。全球每年约有200万新发病例，是女性癌症死亡的主要原因之一。",
                "category": "基础概念"
            },
            {
                "id": 2,
                "title": "乳腺癌症状",
                "content": "乳腺癌常见症状包括：乳房肿块（通常无痛、质地硬、边界不清）、乳头溢液（特别是血性溢液）、乳房皮肤改变（如橘皮样改变、凹陷）、乳头内陷、腋窝淋巴结肿大等。早期乳腺癌可能无明显症状。",
                "category": "症状"
            },
            {
                "id": 3,
                "title": "乳腺癌诊断方法",
                "content": "乳腺癌诊断方法包括：乳腺X线摄影（钼靶）、乳腺超声、乳腺MRI、活检等。细胞学检查是确诊的重要手段。影像学检查可以发现可疑病变，但最终确诊需要病理检查。",
                "category": "诊断"
            },
            {
                "id": 4,
                "title": "乳腺癌分期",
                "content": "乳腺癌分期采用TNM分期系统：T表示肿瘤大小（T1-T4），N表示淋巴结转移（N0-N3），M表示远处转移（M0-M1）。分为0期（原位癌）到IV期（远处转移）。分期越早，预后越好。",
                "category": "分期"
            },
            {
                "id": 5,
                "title": "乳腺癌治疗方法",
                "content": "乳腺癌治疗方法包括：手术治疗（保乳手术、全乳切除术）、化疗、放疗、内分泌治疗、靶向治疗等。治疗方案需根据患者具体情况制定，包括肿瘤分期、分子分型、患者年龄和身体状况。",
                "category": "治疗"
            },
            {
                "id": 6,
                "title": "乳腺癌预防",
                "content": "乳腺癌预防措施包括：定期乳腺检查（自检、临床检查、影像学检查）、保持健康生活方式（规律运动、健康饮食、控制体重）、避免长期使用雌激素、遗传咨询等。高危人群应加强监测。",
                "category": "预防"
            },
            {
                "id": 7,
                "title": "乳腺癌预后",
                "content": "乳腺癌预后与分期、分子分型、治疗反应等因素相关。早期发现和治疗的患者5年生存率可达90%以上。I期患者5年生存率超过95%，IV期患者约为25%。",
                "category": "预后"
            },
            {
                "id": 8,
                "title": "乳腺癌分子分型",
                "content": "乳腺癌分子分型包括：Luminal A型（激素受体阳性，HER2阴性，Ki67低）、Luminal B型（激素受体阳性，HER2阳性或Ki67高）、HER2阳性型、三阴性型（激素受体和HER2均阴性）。不同分型需要不同的治疗方案。",
                "category": "分子分型"
            },
            {
                "id": 9,
                "title": "乳腺癌筛查",
                "content": "乳腺癌筛查包括：乳腺自检（每月一次）、临床乳腺检查（每年一次）、乳腺X线摄影（40岁以上女性每年一次）。高危人群可能需要更频繁的筛查，包括MRI检查。",
                "category": "筛查"
            },
            {
                "id": 10,
                "title": "乳腺癌遗传因素",
                "content": "乳腺癌遗传因素包括BRCA1、BRCA2基因突变等。有家族史的女性应进行遗传咨询和基因检测。BRCA1/2突变携带者一生中患乳腺癌的风险高达60-80%。",
                "category": "遗传"
            },
            {
                "id": 11,
                "title": "乳腺癌危险因素",
                "content": "乳腺癌危险因素包括：年龄（50岁以上风险增加）、家族史、BRCA基因突变、月经初潮早、绝经晚、未生育或晚育、长期使用雌激素、肥胖、饮酒、缺乏运动等。",
                "category": "危险因素"
            },
            {
                "id": 12,
                "title": "乳腺癌术后护理",
                "content": "乳腺癌术后护理包括：伤口护理、引流管管理、功能锻炼、心理支持等。术后应避免患侧上肢过度用力，进行渐进式功能锻炼，预防淋巴水肿。定期复查监测复发。",
                "category": "术后护理"
            },
            {
                "id": 13,
                "title": "乳腺癌化疗",
                "content": "乳腺癌化疗是重要的辅助治疗手段，常用药物包括：蒽环类（阿霉素、表柔比星）、紫杉类（紫杉醇、多西他赛）、环磷酰胺等。化疗方案根据患者情况个体化制定，常见副作用包括脱发、恶心、骨髓抑制等。",
                "category": "化疗"
            },
            {
                "id": 14,
                "title": "乳腺癌内分泌治疗",
                "content": "乳腺癌内分泌治疗适用于激素受体阳性的患者，常用药物包括：他莫昔芬、芳香化酶抑制剂（来曲唑、阿那曲唑、依西美坦）等。治疗周期通常为5-10年，需要定期监测副作用。",
                "category": "内分泌治疗"
            },
            {
                "id": 15,
                "title": "乳腺癌靶向治疗",
                "content": "乳腺癌靶向治疗针对特定分子靶点，包括：HER2靶向药物（曲妥珠单抗、帕妥珠单抗）、CDK4/6抑制剂（哌柏西利、瑞博西利）、PARP抑制剂等。靶向治疗通常与化疗或内分泌治疗联合使用。",
                "category": "靶向治疗"
            },
            {
                "id": 16,
                "title": "乳腺癌免疫治疗",
                "content": "乳腺癌免疫治疗是新兴的治疗方法，通过激活患者自身免疫系统来攻击肿瘤细胞。目前主要用于三阴性乳腺癌的治疗，常用药物包括PD-1/PD-L1抑制剂。",
                "category": "免疫治疗"
            },
            {
                "id": 17,
                "title": "乳腺癌复发转移",
                "content": "乳腺癌复发转移是治疗失败的主要原因，包括局部复发和远处转移。常见转移部位包括：骨、肺、肝、脑等。复发转移的治疗需要多学科协作，包括化疗、放疗、靶向治疗等。",
                "category": "复发转移"
            },
            {
                "id": 18,
                "title": "乳腺癌心理支持",
                "content": "乳腺癌患者常面临巨大的心理压力，需要全面的心理支持。包括：疾病认知教育、情绪管理、家庭支持、社会支持、专业心理咨询等。良好的心理状态有助于提高治疗效果和生活质量。",
                "category": "心理支持"
            },
            {
                "id": 19,
                "title": "乳腺癌营养支持",
                "content": "乳腺癌患者的营养支持非常重要，包括：均衡饮食、适量蛋白质、充足维生素和矿物质、控制体重等。治疗期间应避免过度节食，保证营养充足，有助于提高治疗效果和减少副作用。",
                "category": "营养支持"
            },
            {
                "id": 20,
                "title": "乳腺癌康复锻炼",
                "content": "乳腺癌患者康复锻炼包括：术后功能锻炼、有氧运动、力量训练、柔韧性训练等。适当的运动有助于改善心肺功能、增强免疫力、减少疲劳、提高生活质量。应在医生指导下进行。",
                "category": "康复锻炼"
            }
        ]
        
        self.knowledge_base = default_knowledge
        self._save_knowledge_base()
        self._generate_embeddings()
    
    def _save_knowledge_base(self):
        """保存知识库到文件"""
        try:
            with open(self.knowledge_file, 'w', encoding='utf-8') as f:
                json.dump(self.knowledge_base, f, ensure_ascii=False, indent=2)
            logger.info("知识库保存成功")
        except Exception as e:
            logger.error(f"保存知识库失败: {e}")
    
    def _generate_embeddings(self):
        """生成知识库的向量嵌入"""
        try:
            logger.info("生成知识库向量嵌入...")
            texts = [item['content'] for item in self.knowledge_base]
            self.embeddings = self.embedding_model.encode(texts)
            
            # 保存嵌入向量
            with open(self.embeddings_file, 'wb') as f:
                pickle.dump(self.embeddings, f)
            
            logger.info("向量嵌入生成完成")
        except Exception as e:
            logger.error(f"生成向量嵌入失败: {e}")
            raise
    
    def add_knowledge(self, title: str, content: str, category: str = "其他"):
        """添加新知识到知识库"""
        try:
            new_id = max([item['id'] for item in self.knowledge_base], default=0) + 1
            new_knowledge = {
                "id": new_id,
                "title": title,
                "content": content,
                "category": category
            }
            
            self.knowledge_base.append(new_knowledge)
            
            # 重新生成嵌入向量
            self._generate_embeddings()
            self._save_knowledge_base()
            
            logger.info(f"新知识添加成功: {title}")
            return True
        except Exception as e:
            logger.error(f"添加知识失败: {e}")
            return False
    
    def search(self, query: str, top_k: int = 3) -> List[Dict]:
        """
        搜索相关知识
        
        Args:
            query: 查询文本
            top_k: 返回最相关的前k条知识
            
        Returns:
            相关知识的列表
        """
        try:
            if self.embeddings is None or len(self.embeddings) == 0:
                logger.warning("知识库为空，无法搜索")
                return []
            
            # 生成查询向量
            query_embedding = self.embedding_model.encode([query])
            
            # 计算相似度
            similarities = cosine_similarity(query_embedding, self.embeddings)[0]
            
            # 获取最相关的知识
            top_indices = np.argsort(similarities)[::-1][:top_k]
            
            results = []
            for idx in top_indices:
                if similarities[idx] > 0.3:  # 相似度阈值
                    knowledge = self.knowledge_base[idx].copy()
                    knowledge['similarity'] = float(similarities[idx])
                    results.append(knowledge)
            
            logger.info(f"搜索完成，找到 {len(results)} 条相关知识")
            return results
            
        except Exception as e:
            logger.error(f"搜索失败: {e}")
            return []
    
    def get_answer(self, question: str, context_limit: int = 2) -> Dict:
        """
        基于RAG生成答案
        
        Args:
            question: 用户问题
            context_limit: 使用的上下文知识条数
            
        Returns:
            包含答案和相关知识的字典
        """
        try:
            # 搜索相关知识
            relevant_knowledge = self.search(question, top_k=context_limit)
            
            if not relevant_knowledge:
                return {
                    "answer": "抱歉，我没有找到相关的医学知识来回答您的问题。建议您咨询专业医生。",
                    "sources": [],
                    "confidence": 0.0
                }
            
            # 构建答案
            answer_parts = []
            sources = []
            
            for knowledge in relevant_knowledge:
                answer_parts.append(f"根据医学知识：{knowledge['content']}")
                sources.append({
                    "title": knowledge['title'],
                    "category": knowledge['category'],
                    "similarity": knowledge['similarity']
                })
            
            answer = " ".join(answer_parts)
            
            # 计算置信度（基于最高相似度）
            confidence = max([k['similarity'] for k in relevant_knowledge])
            
            return {
                "answer": answer,
                "sources": sources,
                "confidence": confidence
            }
            
        except Exception as e:
            logger.error(f"生成答案失败: {e}")
            return {
                "answer": "抱歉，系统出现错误，无法生成答案。",
                "sources": [],
                "confidence": 0.0
            }
    
    def get_knowledge_by_category(self, category: str) -> List[Dict]:
        """根据类别获取知识"""
        return [item for item in self.knowledge_base if item['category'] == category]
    
    def get_all_categories(self) -> List[str]:
        """获取所有知识类别"""
        categories = set(item['category'] for item in self.knowledge_base)
        return list(categories)
    
    def update_knowledge(self, knowledge_id: int, title: str = None, content: str = None, category: str = None) -> bool:
        """更新知识库中的知识"""
        try:
            for item in self.knowledge_base:
                if item['id'] == knowledge_id:
                    if title is not None:
                        item['title'] = title
                    if content is not None:
                        item['content'] = content
                    if category is not None:
                        item['category'] = category
                    
                    # 重新生成嵌入向量
                    self._generate_embeddings()
                    self._save_knowledge_base()
                    
                    logger.info(f"知识更新成功: ID {knowledge_id}")
                    return True
            
            logger.warning(f"未找到ID为 {knowledge_id} 的知识")
            return False
        except Exception as e:
            logger.error(f"更新知识失败: {e}")
            return False
    
    def delete_knowledge(self, knowledge_id: int) -> bool:
        """删除知识库中的知识"""
        try:
            original_length = len(self.knowledge_base)
            self.knowledge_base = [item for item in self.knowledge_base if item['id'] != knowledge_id]
            
            if len(self.knowledge_base) < original_length:
                # 重新生成嵌入向量
                self._generate_embeddings()
                self._save_knowledge_base()
                
                logger.info(f"知识删除成功: ID {knowledge_id}")
                return True
            else:
                logger.warning(f"未找到ID为 {knowledge_id} 的知识")
                return False
        except Exception as e:
            logger.error(f"删除知识失败: {e}")
            return False
    
    def get_knowledge_stats(self) -> Dict:
        """获取知识库统计信息"""
        try:
            total_knowledge = len(self.knowledge_base)
            categories = self.get_all_categories()
            category_counts = {}
            
            for item in self.knowledge_base:
                category = item['category']
                category_counts[category] = category_counts.get(category, 0) + 1
            
            return {
                'total_knowledge': total_knowledge,
                'total_categories': len(categories),
                'categories': categories,
                'category_counts': category_counts
            }
        except Exception as e:
            logger.error(f"获取知识库统计信息失败: {e}")
            return {
                'total_knowledge': 0,
                'total_categories': 0,
                'categories': [],
                'category_counts': {}
            }
